package com.sensei.search.nodes;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.Map.Entry;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

import org.apache.log4j.Logger;

import com.browseengine.bobo.api.BrowseRequest;
import com.browseengine.bobo.api.BrowseResult;
import com.linkedin.nimbus.NimbusException;
import com.linkedin.nimbus.cluster.Cluster;
import com.linkedin.nimbus.cluster.ClusterListenerAdapter;
import com.linkedin.nimbus.cluster.Node;
import com.linkedin.nimbus.cluster.Router;
import com.sensei.search.client.ResultMerger;
import com.sensei.search.client.SenseiClientHandler;
import com.sensei.search.client.SenseiConnResource;
import com.sensei.search.client.SenseiConnResourceFactory;
import com.sensei.search.client.SenseiConnResource.HostAddr;
import com.sensei.search.util.pool.KeyedResourcePool;
import com.sensei.search.util.pool.ResourcePoolConfig;

public class SenseiBroker extends ClusterListenerAdapter  {
	private static Logger logger = Logger.getLogger(SenseiBroker.class);
	
	private final Cluster _cluster;
	private ExecutorService _threadPool = Executors.newCachedThreadPool();
	
	private volatile Map<Integer,ArrayList<Node>> _partToNodeMap;
	
	private final Random _rand;
	
	private final KeyedResourcePool<HostAddr,SenseiConnResource> _connectionPool;
	
	private static Callable<BrowseResult> newCallable(final BrowseRequest req,final Node node,final KeyedResourcePool<HostAddr,SenseiConnResource> connPool){
	    return new Callable<BrowseResult>(){
	            public BrowseResult call() throws Exception {
	            	SenseiConnResource conn = null;
	            	HostAddr addr = new HostAddr(node.host().getHostAddress(),node.port());
	            	try{
	                  conn = connPool.checkout(addr);
	                  SenseiClientHandler svc = conn.getClientHandler();
					  return svc.browse(req);
					}
					finally{
					  if (conn!=null){
						  connPool.checkin(addr, conn);
					  }
					}
	            }
	            
	    };
	}
	
	private static BrowseResult doBrowse(ExecutorService threadPool,BrowseRequest req,Node[] nodes,KeyedResourcePool<HostAddr,SenseiConnResource> connectionPool){
		if (nodes!=null && nodes.length>0){
		  int oldOffset = req.getOffset();
		  int oldCount = req.getCount();
		  
		  req.setOffset(0);
		  req.setCount(oldCount);
		  
	      List<Future<BrowseResult>> resultFutures = new ArrayList<Future<BrowseResult>>(nodes.length);
	      
	      List<BrowseResult> resultList = new ArrayList<BrowseResult>(nodes.length);
	      for (Node node : nodes){
	    	  Callable<BrowseResult> callable = newCallable(req,node,connectionPool);
	    	  resultFutures.add(threadPool.submit(callable));
	      }
	      
	      for (Future<BrowseResult> resultFuture : resultFutures){
	    	  try {
				resultList.add(resultFuture.get());
			  } catch (InterruptedException e) {
				logger.error(e.getMessage(),e);
			  } catch (ExecutionException e) {
				logger.error(e.getMessage(),e);
			  }
	      }
	      
	      BrowseResult res = ResultMerger.merge(req, resultList);
	      req.setOffset(oldOffset);
	      req.setCount(oldCount);
	      return res;
		}
		else{
			logger.warn("no server exist to handle request.");
			return new BrowseResult();
		}
	}
	
	public SenseiBroker(Cluster cluster) throws NimbusException{
		_cluster = cluster;
		_cluster.addListener(this);
		_rand = new Random();
		
		ResourcePoolConfig poolConfig = new ResourcePoolConfig();
		poolConfig.setIsFair(false);
		poolConfig.setMaxInvalidAttempts(5);
		poolConfig.setMaxPoolSize(100);
		poolConfig.setTimeout(5000, TimeUnit.MILLISECONDS);
		
		_connectionPool = new KeyedResourcePool<HostAddr, SenseiConnResource>(new SenseiConnResourceFactory(), poolConfig);
		
		/**
		Router router = _cluster.router();
		 * hmm.. this does not work... 
		if (router!=null){
			System.out.println("count: "+cluster.nodes().length);
			applyPartitionMap(router.partitionToNodeMap());
		}
		*/
		
		applyNodes(cluster.nodes());
	}
	
	public BrowseResult browse(BrowseRequest req){
		Node[] nodes = null;
		Map<Integer,ArrayList<Node>> map = _partToNodeMap;
		if (map!=null && map.size() > 0){
			ArrayList<Node> nodeArray = new ArrayList<Node>(map.size());
			Collection<ArrayList<Node>> nodeColl = map.values();
			for (ArrayList<Node> reps : nodeColl){
				int size = reps.size();
				if (size>0){
					if (reps.size() > 1){
						int idx = _rand.nextInt(size);
						nodeArray.add(reps.get(idx));
					}
					else{
						nodeArray.add(reps.get(0));
					}
				}
			}
			nodes = nodeArray.toArray(new Node[nodeArray.size()]);
		}
		return doBrowse(_threadPool,req,nodes,_connectionPool);
	}
	
	private void applyNodes(Node[] nodes){
		HashMap<Integer,ArrayList<Node>> partToNodeMap = new HashMap<Integer,ArrayList<Node>>();
	    for (Node node : nodes){
	    	  int[] parts = node.partitions();
	    	  for (int part : parts){
	    		  ArrayList<Node> nodeList = partToNodeMap.get(part);
	    		  if (nodeList == null){
	    			  nodeList = new ArrayList<Node>();
	    			  partToNodeMap.put(part, nodeList);
	    		  }
	    		  nodeList.add(node);
	    	  }
	    }
	    
	    _partToNodeMap = partToNodeMap; 
	}
	
	private void applyPartitionMap(Map<Integer,Set<Node>> partToNodeMap){
		System.out.println("partmap :"+partToNodeMap);
		if (partToNodeMap!=null){
			HashMap<Integer,ArrayList<Node>> map = new HashMap<Integer,ArrayList<Node>>();
			Set<Entry<Integer,Set<Node>>> entries = partToNodeMap.entrySet();
			for (Entry<Integer,Set<Node>> entry : entries){
				Set<Node> nodeSet = entry.getValue();
				ArrayList<Node> nodeList = new ArrayList<Node>(nodeSet.size());
				nodeList.addAll(nodeSet);
				map.put(entry.getKey(), nodeList);
			}
			_partToNodeMap = map;
		}
	}
	
	@Override
	public void handleClusterConnected(Node[] nodes, Router route) {
		System.out.println("handleClusterConnected called");
		//applyPartitionMap(route.partitionToNodeMap());
		applyNodes(nodes);
	}


	@Override
	public void handleClusterNodesChanged(Node[] nodes, Router route) {
		System.out.println("handleClusterNodeChanged called");
		//applyPartitionMap(route.partitionToNodeMap());
		applyNodes(nodes);
	}

	@Override
	public void handleClusterSessionExpired() {
		logger.warn("Cluster was shut down, reconnecting...");
		while(true){
			  try{
			    _cluster.awaitConnection();
			    break;
			  }
			  catch(InterruptedException ie){
				  continue;
			  }
		  }
		  logger.info("Reconnected to cluster");
	}
	
	public void shutdown(){
		this._threadPool.shutdown();
		this._connectionPool.close();
		this._partToNodeMap.clear();
	}
}
